# 导入第三方库（pip可下载）
import matplotlib.pyplot
import numpy

# 导入自定义库
from dapas.statistics import *
from dapas.distribution import *


# 套壳函数（方便换函数）
def f(x):
    y = geom(x, 0.2)     # 更换这里的函数即可
    return y


# 生成数据与概率密度计算
X = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
Y = []
for x_i in X:
    Y.append(f(x_i))

Y_0 = []
for x_i in X:
    Y_0.append(hyper(x_i, 50, 12, 100))

Y_1 = []
for x_i in X:
    Y_1.append(binom(x_i, 0.2, 20))

Y_2 = []
for x_i in X:
    Y_2.append(pois(x_i, 4))


# 绘制图像
# 图1
matplotlib.pyplot.figure(num=1, figsize=(10, 5))
matplotlib.pyplot.scatter(X, Y, color="blue", label="Geometry Distribution (p=0.2)")   # 散点图
matplotlib.pyplot.scatter(X, Y_0, color="red", label="Hyper-Geometry Distribution (n=50, M=12, N=100)")   # 散点图
matplotlib.pyplot.legend(loc="upper right")     # 显示图例

# 图1：四张图均匀分布
matplotlib.pyplot.figure(num=2, figsize=(10, 8))
# 子图1
matplotlib.pyplot.subplot(2, 1, 1)
matplotlib.pyplot.bar(X, Y_1, width=0.5, color="blue", label="Binomial Distribution (p=0.2, n=20)")   # 条形图
matplotlib.pyplot.legend(loc="upper right")     # 显示图例
# 子图2
matplotlib.pyplot.subplot(2, 1, 2)
matplotlib.pyplot.bar(X, Y_2, width=0.5, color="red", label="Poison Distribution (λ = np = 4)")   # 条形图
matplotlib.pyplot.legend(loc="upper right")     # 显示图例


# 显示绘图结果
matplotlib.pyplot.show()
